// Copyright 2020 The SQLFlow Authors. All rights reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package step

import (
	_ "image/png"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
	"sqlflow.org/sqlflow/go/database"
	pb "sqlflow.org/sqlflow/go/proto"
	"sqlflow.org/sqlflow/go/step/tablewriter"
)

var testImageHTML string = "<div align='center'><img src='' /></div>"

func makeTestSession(dbConnStr string) *pb.Session {
	return &pb.Session{DbConnStr: dbConnStr}
}

func TestStepTrainSQL(t *testing.T) {
	if os.Getenv("SQLFLOW_TEST_DB") != "mysql" {
		t.Skip("skip no mysql test.")
	}
	a := assert.New(t)
	session := makeTestSession(database.GetTestingMySQLURL())

	sql := `SELECT * FROM iris.train WHERE class!=2
	TO TRAIN DNNClassifier
	WITH
		model.n_classes = 2,
		model.hidden_units = [10, 10],
		train.batch_size = 4,
		validation.select = "SELECT * FROM iris.test WHERE class!=2",
		validation.metrics = "Accuracy,AUC"
	LABEL class
	INTO sqlflow_models.mytest_model;`
	table, e := tablewriter.Create("ascii", 100, os.Stdout)
	a.NoError(e)
	out, e := GetStdout(func() error { return RunSQLProgramAndPrintResult(sql, session, table, false, false) })
	a.NoError(e)
	a.Contains(out, "Done training")
}
func TestImage(t *testing.T) {
	a := assert.New(t)
	a.True(isHTMLCode("<div></div>"))
	_, e := getBase64EncodedImage("")
	a.Error(e)
	image, e := getBase64EncodedImage(testImageHTML)
	a.NoError(e)
	a.NoError(imageCat(image)) // sixel mode
}
